package io.github.junxworks.junx.stream.extractor.db;

import com.google.common.collect.Maps;
import io.github.junxworks.junx.stream.extractor.AbstractExtractor;
import io.github.junxworks.junx.stream.extractor.ExtractContext;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.util.Map;

/**
 * DB数据抽取器，只从数据库中抽取数据
 *
 * @author 王兴
 * @date 2024/09/10
 */
@Slf4j
@Setter
@Getter
public class DBDataExtractor extends AbstractExtractor {

    private DataSource dataSource;

    @Override
    public void extract(ExtractContext context) throws Exception {
        super.beforeExtract(context);
        int batchSize = context.getBatchSize();
        try (Connection conn = dataSource.getConnection(); PreparedStatement ps = conn.prepareStatement(context.getSql(), ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);) {
            conn.setAutoCommit(false);
            ps.setFetchSize(batchSize);
            log.info(this.flowName + "数据正在查询中......");
            clock.reset();
            try (ResultSet rs = ps.executeQuery();) {
                log.info(this.flowName + "数据查询完毕，耗时:" + clock.countMillis() + " ms.");
                ResultSetMetaData md = rs.getMetaData();
                int columnCount = md.getColumnCount();
                int currentSize = 0;
                while (rs.next()) {
                    Map<String, Object> map = Maps.newHashMap();
                    for (int i = 1; i <= columnCount; i++) {
                        String columnName = md.getColumnName(i);
                        map.put(columnName, rs.getObject(i));
                    }
                    if (convertor != null) {
                        dataBoat.add(convertor.convert(map));
                    } else {
                        dataBoat.add(map);
                    }
                    if (++currentSize == batchSize) {
                        currentSize = 0;
                        flush();
                    }
                }
                if (!dataBoat.isEmpty()) {
                    flush();
                }
            }
        } finally {
            this.stop();
        }
    }

}
